import numpy as np
import math
import matplotlib.pyplot as plt


class SysParam:
    def __init__(self):
        pass

class SimParam:
    def __init__(self):
        pass

class Channel:
    def __init__(self,param):
        self.sigm = math.sqrt(param.n0/2.)
        self.df = param.FreqErr/param.FreqCarrier*param.Rate
    
    def awgn(self,x):
        n = np.random.randn(len(x))+1j*np.random.randn(len(x))
        n *= self.sigm
        return x+n

    def doit(self,x):
        return self.awgn(self.ferr(x))

    def ferr(self,x):
        phase = np.arange( len(x) )*self.df
        y = np.exp( 1j * phase)*x
        return y

class DBPSK:
    def __init__(self,param):
        self.tPhase = 0.

    def modu(self,x):
        res = [self.tPhase]
        for d in x:
            if d == 1:
                self.tPhase += np.pi
                self.tPhase %= 2*np.pi
            res.append(self.tPhase)
        return np.exp(-1j*np.array(res))

    def demodu(self,x):
        r = x[1:] * np.conj(x[:-1])
        return np.array([ 0 if x>0. else 1 for x in np.real(r)])

class BPSK:
    def __init__(self,param):
        pass

    def modu(self,x):
        res = x*np.pi
        return np.exp(-1j*np.array(res))

    def demodu(self,x):
        r = x
        return np.array([ 0 if x>0. else 1 for x in np.real(r)])

class Source:
    def __init__(self,len):
        self.d = np.random.randint(0,2,(len,))
    
    def error(self,y):
        return sum(np.bitwise_xor(self.d,y))

def sim(Sys,simP,sysP):
    
    sysP.n0 = math.pow(10,-simP.EbN0/10.)

    sys = Sys(sysP)

    ch = Channel(sysP)

    s = Source(simP.len)

    if simP.verbose:
        print("d:",s.d)

    t = sys.modu(s.d)
    if simP.verbose:
        print("t:",t)

    r = ch.doit(t)
    if simP.verbose:
        print("r:",r)

    y = sys.demodu(r)
    if simP.verbose:
        print("y:",y)

    err = s.error(y)

    # print(err)
    
    return err/simP.len

def main():
    p = SysParam()
    p.Rate = 1.0
    p.FreqCarrier = 1e9
    p.FreqErr = 0
    
    sp = SimParam()
    sp.EbN0 = 5
    sp.len = 100
    sp.verbose = True
    sp.ebn0s = np.arange(6,10,0.1)
    
    sim(BPSK,sp,p)
    
    rd = []
    rb = []
    sp.len = 1000000
    sp.verbose = False

    for EbN0 in sp.ebn0s:
        sp.EbN0 = EbN0
        e = sim(BPSK,sp,p)
        print("BPSK",EbN0,e)
        rb.append(e)
        e = sim(DBPSK,sp,p)
        print("DBPSK",EbN0,e)
        rd.append(e)
    import matplotlib.pyplot as plt

    fig = plt.figure(figsize=(8, 8)) # You were missing the =
    ax = fig.add_subplot(1, 1, 1)
    ax.set_yscale('log')
    plt.grid(True)
    plt.plot(sp.ebn0s,np.array(rb),'r')
    plt.plot(sp.ebn0s,np.array(rd),'b')
    plt.ylim(0.000001, 0.01) # Fix the x limits to fit all the points
    plt.show()

if __name__ == "__main__":
    main()
                
